package com.zengde.ai_sandbox.controller;

import com.zengde.ai_sandbox.service.QdrantService;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.*;
import java.util.List;
import java.util.Map;
import java.util.HashMap;

@RestController
@RequestMapping("/api/vector")
public class QdrantController {

    @Autowired
    private QdrantService qdrantService;

    @PostMapping("/add")
    public ResponseEntity<?> addVector(
            @RequestParam Long id,
            @RequestBody List<Float> vector) {
        try {
            // 检查服务是否可用
            if (!qdrantService.isQdrantAvailable()) {
                return ResponseEntity.internalServerError().body(Map.of(
                    "error", "Qdrant服务不可用，请确保Qdrant已启动在http://localhost:6334"
                ));
            }
            
            qdrantService.addVector(id, vector);
            return ResponseEntity.ok(Map.of(
                "message", "Vector added successfully",
                "id", id
            ));
        } catch (Exception e) {
            return ResponseEntity.badRequest().body(Map.of(
                "error", e.getMessage()
            ));
        }
    }

    @GetMapping("/search")
    public ResponseEntity<?> searchSimilar(
            @RequestBody List<Float> vector,
            @RequestParam(defaultValue = "5") int limit) {
        try {
            // 检查服务是否可用
            if (!qdrantService.isQdrantAvailable()) {
                return ResponseEntity.internalServerError().body(Map.of(
                    "error", "Qdrant服务不可用，请确保Qdrant已启动在http://localhost:6334"
                ));
            }
            
            List<Map<String, Object>> results = qdrantService.searchSimilar(vector, limit);
            
            Map<String, Object> response = new HashMap<>();
            response.put("query", vector);
            response.put("results", results);
            
            return ResponseEntity.ok(response);
        } catch (Exception e) {
            return ResponseEntity.badRequest().body(Map.of(
                "error", e.getMessage()
            ));
        }
    }

    @PostMapping("/collection")
    public ResponseEntity<?> createCollection() {
        try {
            // 检查服务是否可用
            if (!qdrantService.isQdrantAvailable()) {
                return ResponseEntity.internalServerError().body(Map.of(
                    "error", "Qdrant服务不可用，请确保Qdrant已启动在http://localhost:6334"
                ));
            }
            
            qdrantService.createCollection();
            return ResponseEntity.ok(Map.of(
                "message", "Collection created successfully"
            ));
        } catch (Exception e) {
            return ResponseEntity.badRequest().body(Map.of(
                "error", e.getMessage()
            ));
        }
    }
    
    @GetMapping("/status")
    public ResponseEntity<?> checkStatus() {
        boolean isAvailable = qdrantService.isQdrantAvailable();
        
        Map<String, Object> response = new HashMap<>();
        response.put("service", "qdrant");
        response.put("available", isAvailable);
        
        if (!isAvailable) {
            response.put("message", "Qdrant服务不可用，请确保Qdrant已启动在http://localhost:6334");
        }
        
        return ResponseEntity.ok(response);
    }
}
